#

import heapq
from collections import defaultdict
import os
import json
from settings import config


class ProxyDB(object):
    def __init__(self, db=None, save_path=None):
        self.db = defaultdict(set)
        if db and os.path.exists(db):
            data = json.load(open(db, 'r', encoding='utf8'))
            for k in data:
                self.db[k] = set([tuple(v) for v in data[k]])
        self.save_path = save_path

    def zadd(self, key, value):
        if self._valid(value):
            if not self._exists(key, value):
                self.db[key].add(value)
                if self._exists(key, value):
                    return value

    def zrem(self, key, value):
        if self._exists(key, value):
            self.db[key].remove(value)
            return True

    def zpop(self, key):
        keys = self.db[key]
        heapq.heapify(keys)
        rv = heapq.nlargest(1, keys, lambda s: s[1])
        self.db[key].remove(rv)
        return rv

    def zpopleft(self, key):
        keys = self.db[key]
        heapq.heapify(keys)
        return heapq.heappop(keys)

    def zrange(self, key, start, stop):
        if self.zcount(key) < 1:
            return set()
        keys = self.db[key]
        score = [k[1] for k in keys]
        MIN = min(start, min(score))
        MAX = max(stop, max(score))
        rv = [k for k in keys if k[1] >= MIN and k[1] <= MAX]
        return rv

    def zsmalls(self, key, number):
        keys = list(self.db[key])
        heapq.heapify(keys)
        rv = heapq.nsmallest(number, keys, key=lambda s: s[1])
        return rv

    def zlarges(self, key, number):
        keys = list(self.db[key])
        heapq.heapify(keys)
        rv = heapq.nlargest(number, keys, key=lambda s: s[1])
        return rv

    def zall(self, key):
        return list(self.db[key])

    def zcount(self, key):
        return len(self.db[key])

    def zmax(self, key, value):
        if self.zrem(key, value):
            value = list(value)
            value[1] = 100
            value = tuple(value)
            return self.zadd(key, value)

    def zdecrease(self, key, value):
        if self.zrem(key, value):
            value = list(value)
            value[1] -= 1
            value = tuple(value)
            return self.zadd(key, value)

    def zflush(self, key, score=config['flush_score']):
        proxy_flush = self.zrange(key, 0, score)
        if not proxy_flush:
            return
        for proxy in proxy_flush:
            self.zrem(key, proxy)

    def _valid(self, value):
        if not isinstance(value, (tuple, list)):
            raise ValueError('value must be tuple or list')
        if len(value) != 2:
            raise ValueError('value length must be 2')
        if not isinstance(value[0], (str,)):
            raise ValueError('value[0] must be str')
        if not isinstance(value[1], (int, float)):
            raise ValueError('value[1] must be int or float')
        return True

    def _exists(self, key, value):
        if value in self.db[key]:
            return True

    def save(self):
        fname = f'db_temp.json'
        if self.save_path:
            path = os.path.join(self.save_path, fname)
            rename_path = os.path.join(self.save_path, 'db.json')
        else:
            path = fname
            rename_path = 'db.json'
        data = {}
        for k in self.db:
            data[k] = list(self.db[k])
        json.dump(data, open(path, 'w', encoding='utf8'))
        if os.path.exists(rename_path):
            os.remove(rename_path)
        os.rename(path, rename_path)
